//
// Created by WJY on 2021/8/8.
//

#ifndef CVM_INCLUDE_CVM_PACKED_FUNCTION_H_
#define CVM_INCLUDE_CVM_PACKED_FUNCTION_H_

#include <functional>

namespace tvm {
namespace runtime {

class TVMValue;
class TVMArgs;
class TVMRetValue;

class PackedFunc {
 public:
  using FType = std::function<void(TVMArgs args, TVMRetValue* rv)>;

  PackedFunc() = default;
  PackedFunc(std::nullptr_t null) {}  // NOLINT
  explicit PackedFunc(FType body) : body_(body) {}

  template <typename... Args>
  inline TVMRetValue operator()(Args&&... args) const;

  inline void CallPacked(TVMArgs args, TVMRetValue* rv) const;

  inline FType body() const;

  bool operator==(std::nullptr_t null) { return body_ == nullptr; }
  bool operator!=(std::nullptr_t null) { return body_ != nullptr; }

 private:
  FType body_;
};

class TVMArgs {
 public:
  TVMArgs(TVMValue* values, const int* type_codes, int num_args)
      : values_(values), type_codes_(type_codes), num_args_(num_args) {}

  inline int size() const;

  inline TVMValue operator[](int i) const;

 private:
  const TVMValue* values_;
  const int* type_codes_;
  int num_args_;
};

}  // namespace runtime
}  // namespace tvm

#endif  // CVM_INCLUDE_CVM_PACKED_FUNCTION_H_
